import json
from multiprocessing.connection import Listener
import time
from threading import Thread, Condition
from repeating_timer import RepeatingTimer

class Scheduler:
    def __init__(self, agent_colors: list, ip:str, ports, authkey):
        self._agent_num = len(agent_colors)  # 智能体数量
        self._agent_colors = agent_colors  # 智能体所属方
        self._ip = ip
        if isinstance(ports, list):  # ports支持输入list或int两种类型，此处转换成list
            self._ports = ports
        elif isinstance(ports, int):
            self._ports = list(range(ports, ports+self._agent_num+1))
        else:
            raise TypeError("Port must be int or list")
        self._authkey = authkey
        self._listeners = []
        for i in range(self._agent_num+1):
            self._listeners.append(Listener(address=(self._ip, self._ports[i]), authkey=self._authkey))
        self._is_ready = [False] * (self._agent_num+1)
        self.observations = ""  # 所有的态势信息组成的字符串
        self.observation_dict = []  # 按照不同对抗方，如"Red", "Blue", "All"等，区分的态势信息
        self.actions = ""
        self._cond_agent_start = Condition()  # Agent推演阶段开始的同步
        self._cond_agent_end = Condition()  # Agent推演阶段结束的同步
        self._cond_engine_end = Condition()  # Game Evolver推演阶段结束的同步
        self._game_over = False
        self.stage = "starting"  # 游戏的三个阶段：starting, gaming, ending

    def start(self):
        # === starting 阶段 ===
        print(f"\t=== {self.stage} 阶段 ===")
        self._timer = RepeatingTimer(0.5, self._cond_agent_start, self._cond_agent_end, self._cond_engine_end)
        transfer_engine_thread = Thread(target=self.transfer_engine, daemon=True)
        for i in range(1, self._agent_num+1):
            transfer_agent_thread = Thread(target=self.transfer_agent, args=(i,), daemon=True)
            transfer_agent_thread.start()
        transfer_engine_thread.start()
        while not all(self._is_ready):
            pass
        # === gaming 阶段 ===
        self.stage = "gaming"
        print(f"\t=== {self.stage} 阶段 ===")
        self._timer.start()
        self._timer.join()

    def transfer_engine(self):
        conn_engine = self._listeners[0].accept()
        print("engine连接上了")
        print("\t  Scheduler -> Engine: 'ready'")
        conn_engine.send("ready")
        recv_data = conn_engine.recv()
        print(f"\t  Scheduler <- Engine: {recv_data}")
        if recv_data == "ok": # 接收到准备好的信号
            self._is_ready[0] = True
            while not all(self._is_ready):  # 全部准备好后，跳出该循环
                pass
        else:
            raise Exception("engine not ready")
        # with self._cond_agent_end or self._cond_engine_end:
        while True:
            with self._cond_agent_end:
                self._cond_agent_end.wait()
            print(f"\t  Scheduler -> Engine: {self.actions}")
            conn_engine.send(self.actions)
            self.actions = ""
            recv_data = conn_engine.recv()
            print(f"\t  Scheduler <- Engine: {recv_data}")
            self.observations = recv_data
            self.split_observations()
            if(self.observation_dict['Done']):  # 接收到了推演"结束"消息
                self._game_over = True
                with self._cond_engine_end:
                    self._cond_engine_end.notify_all()
                print("\t  Scheduler -> Engine: 'done'")
                conn_engine.send("done")  # 告诉game_engine游戏结束
                self.stage = "ending"
                print(f"\t=== {self.stage} 阶段 ===")
                break
            else:
                with self._cond_engine_end:
                    self._cond_engine_end.notify_all()

    def transfer_agent(self, number: int):
        '''
        number: 智能体编号，1,2,...，0号是game_engine
        '''
        conn_agent = self._listeners[number].accept()
        print("agent连接上了")
        print("\t  Scheduler -> Agent: 'ready'")
        conn_agent.send("ready")
        recv_data = conn_agent.recv()
        print(f"\t  Scheduler <- Agent: {recv_data}")
        if recv_data == "ok": # 接收到准备好的信号
            self._is_ready[number] = True
            while not all(self._is_ready):  # 全部准备好后，跳出该循环
                pass
        else:
            raise Exception("engine not ready")
        while True:
            if self._game_over:
                print("\t  Scheduler -> Agent: 'done'")
                conn_agent.send("done")
                self._timer.stop()
                break
            with self._cond_agent_start:
                self._cond_agent_start.wait()
            print(f"\t  Scheduler -> Agent: {self.observation_dict[self.agent_colors[number-1]]}")
            conn_agent.send(self.observation_dict[self.agent_colors[number-1]])
            time.sleep(0.1)
            recv_data = conn_agent.recv()
            print(f"\t  Scheduler <- Agent: {recv_data}")
            self.actions += recv_data

    def split_observations(self):
        '''
        将observations分割成多个智能体的态势
        :param observations: 所有智能体的态势
        :return: 各智能体的态势，dict(str: str)
        '''
        key_values = json.loads(self.observations)
        self.observation_dict = {}
        for key, value in key_values.items():
            if key != "Done":
                self.observation_dict[key] = str(value).replace("'", '"')
